#!/usr/bin/env python
# -*- coding: utf-8 -*-
# -------------------------------------------------------------------------------
# @Time    : 2025/9/28 10:41
# @Author  : WXY
# @File    : YouTubeDownloadWorker.py
# @PROJECT_NAME: youtubedown_gui
# @PRODUCT_NAME: PyCharm
import re
import yt_dlp
import threading
import time
import queue
import sys
import os
from PySide6.QtCore import QThread, Signal, QObject
from LoggerManager import logger_manager

MAX_THREADS = 3
MAX_RETRIES = 3


class YouTubeDownloadWorker(QThread):
    """YouTube下载工作线程"""
    # 信号定义
    progress_updated = Signal(int)  # 进度更新信号 (0-100)
    status_updated = Signal(str)  # 状态更新信号
    download_finished = Signal(bool, str)  # 下载完成信号 (成功/失败, 消息)
    log_message = Signal(str)  # 日志消息信号
    video_info_extracted = Signal(str)  # 视频信息提取信号 (视频标题)

    def __init__(self, url, download_path, resolution=None, only_mp4=False):
        super().__init__()
        self.url = url
        self.download_path = download_path
        self.resolution = self._parse_resolution(resolution)  # 转换分辨率字符串为数字
        self.only_mp4 = only_mp4
        self.is_cancelled = False

        # 下载统计
        self.total_videos = 0
        self.downloaded_videos = 0
        self.current_video_progress = 0
    def _parse_resolution(self, resolution_str):
        """将分辨率字符串转换为数字"""
        if not resolution_str or resolution_str == "不限制":
            return None

        # 分辨率映射表
        resolution_map = {
            "360P": 360,
            "720P": 720,
            "1080P": 1080,
            "2K": 1440,  # 2K 通常指 1440p
            "4K": 2160  # 4K 指 2160p
        }

        return resolution_map.get(resolution_str, None)

    def _check_resolution_available(self, info):
        """检查是否有符合分辨率要求的视频格式"""
        if not self.resolution:
            return True

        formats = info.get('formats', [])
        if not formats:
            return False

        # 查找符合分辨率要求的格式（分辨率 >= 设置值）
        suitable_formats = []

        for fmt in formats:
            height = fmt.get('height')
            if height and height >= self.resolution:  # 寻找大于等于设置分辨率的格式
                # 检查是否有视频流
                vcodec = fmt.get('vcodec', 'none')
                if vcodec != 'none':
                    suitable_formats.append(fmt)

        if suitable_formats:
            self.log_message.emit(f"✅ 找到 {len(suitable_formats)} 个符合分辨率要求的格式")
            # 显示最佳匹配的分辨率（选择最高分辨率的）
            best_format = max(suitable_formats, key=lambda x: x.get('height', 0))
            best_height = best_format.get('height', 0)
            self.log_message.emit(f"📺 将下载最优分辨率: {best_height}p")
            return True
        else:
            # 显示可用的分辨率
            available_heights = []
            for fmt in formats:
                height = fmt.get('height')
                vcodec = fmt.get('vcodec', 'none')
                if height and vcodec != 'none':
                    available_heights.append(height)

            if available_heights:
                available_heights = sorted(set(available_heights), reverse=True)
                max_available = max(available_heights)
                self.log_message.emit(f"📊 可用分辨率: {', '.join(map(str, available_heights))}p")
                self.log_message.emit(f"❌ 当前视频最高分辨率是 {max_available}p，低于要求的 {self.resolution}p")
                self.log_message.emit("💡 建议: 试试降低分辨率限制或选择'不限制'")

            return False

    def sanitize_filename(self, name):
        """清理文件名中的非法字符"""
        return re.sub(r'[\\/:*?"<>|]', '_', name)

    def progress_hook(self, d):
        """yt-dlp进度回调函数"""
        if self.is_cancelled:
            # 抛出异常来强制中断下载
            raise Exception("下载已被用户取消")

        if d['status'] == 'downloading':
            # 尝试获取进度百分比，支持多种数据格式
            percent = None
            total_bytes = None
            downloaded_bytes = None

            # 方法1: 直接从字典获取百分比
            if '_percent_str' in d:
                try:
                    percent_str = d['_percent_str'].strip().replace('%', '')
                    percent = float(percent_str)
                except (ValueError, AttributeError):
                    pass

            # 方法2: 从 downloaded_bytes 和 total_bytes 计算
            if percent is None:
                if 'total_bytes' in d and d['total_bytes'] and 'downloaded_bytes' in d:
                    total_bytes = d['total_bytes']
                    downloaded_bytes = d['downloaded_bytes']
                    percent = (downloaded_bytes / total_bytes) * 100
                elif 'total_bytes_estimate' in d and d['total_bytes_estimate'] and 'downloaded_bytes' in d:
                    total_bytes = d['total_bytes_estimate']
                    downloaded_bytes = d['downloaded_bytes']
                    percent = (downloaded_bytes / total_bytes) * 100

            # 方法3: 从分片信息计算（针对HLS下载）
            if percent is None and 'fragment_index' in d and 'fragment_count' in d:
                if d['fragment_count'] and d['fragment_count'] > 0:
                    percent = (d['fragment_index'] / d['fragment_count']) * 100

            # 如果成功获取到进度百分比
            if percent is not None:
                percent = int(percent)
                self.current_video_progress = percent

                # 计算总体进度
                if self.total_videos > 0:
                    video_progress = (self.downloaded_videos * 100 + percent) / self.total_videos
                    self.progress_updated.emit(int(video_progress))
                else:
                    self.progress_updated.emit(percent)

                # 获取速度信息
                speed = d.get('speed', 0)
                speed_str = ""
                if speed and speed > 0:
                    if speed > 1024 * 1024:
                        speed_str = f" - {speed / 1024 / 1024:.1f} MB/s"
                    else:
                        speed_str = f" - {speed / 1024:.1f} KB/s"

                # 获取文件大小信息
                size_info = ""
                if total_bytes:
                    if total_bytes > 1024 * 1024 * 1024:
                        size_info = f" of {total_bytes / 1024 / 1024 / 1024:.1f} GB"
                    elif total_bytes > 1024 * 1024:
                        size_info = f" of {total_bytes / 1024 / 1024:.1f} MB"
                    else:
                        size_info = f" of {total_bytes / 1024:.1f} KB"

                # 发送状态更新
                status_msg = f"下载中... {percent}%{size_info}{speed_str}"
                self.status_updated.emit(status_msg)

                # 同时发送详细的日志信息
                if 'fragment_index' in d and 'fragment_count' in d:
                    self.log_message.emit(f"📊 {status_msg} (分片 {d['fragment_index']}/{d['fragment_count']})")
                else:
                    self.log_message.emit(f"📊 {status_msg}")
            else:
                # 如果无法获取进度，至少显示正在下载
                self.log_message.emit("📊 正在下载...")

        elif d['status'] == 'finished':
            self.log_message.emit(f"✅ 下载完成: {os.path.basename(d['filename'])}")
            self.log_message.emit("🔄 正在进行后续的合并处理...")  # 添加后处理提示
            self.downloaded_videos += 1

        elif d['status'] == 'error':
            self.log_message.emit(f"❌ 下载错误: {d.get('error', '未知错误')}")

    def get_ydl_opts(self, output_folder=None):
        """生成 yt-dlp 下载配置"""
        if self.resolution:
            format_selector = (
                f"best[height>={self.resolution}]/"  # 选择符合分辨率要求的最佳格式
                f"bestvideo[height>={self.resolution}]+bestaudio/"
                f"best[ext=mp4][height>={self.resolution}]/"
                f"bestvideo[ext=mp4][height>={self.resolution}]+bestaudio[ext=m4a]"
            )
            self.log_message.emit(f"🎯 分辨率限制: 最低 {self.resolution}p（将选择最优格式）")
        else:
            format_selector = (
                "best/"
                "bestvideo+bestaudio/"
                "best[ext=mp4]/"
                "bestvideo[ext=mp4]+bestaudio[ext=m4a]"
            )
            self.log_message.emit("🎯 分辨率限制: 无限制")

        self.log_message.emit(f"📋 格式选择器: {format_selector}")

        outtmpl = '%(title)s.%(ext)s' if output_folder is None else os.path.join(output_folder, '%(title)s.%(ext)s')

        return {
            'http_headers': {
                'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36'
            },
            'socket_timeout': 60,
            'retries': 5,
            'cookies': 'cookies.txt',
            'format': format_selector,
            'merge_output_format': 'mp4',
            'outtmpl': outtmpl,
            'quiet': False,  # False可以获取更多调试信息
            'noplaylist': True,
            'concurrent_fragment_downloads': 5,
            'progress_hooks': [self.progress_hook],  # 添加进度回调
            'postprocessors': [
                {
                    'key': 'FFmpegVideoConvertor',
                    'preferedformat': 'mp4',
                },
                {
                    'key': 'FFmpegMetadata',
                }
            ],
            'postprocessor_args': [
                '-c:v', 'copy',  # 视频不重新编码
                '-c:a', 'aac',  # 音频转为 mp4 兼容格式
            ],
        }

    def download_single_video(self, url, output_folder=None):
        """下载单个视频"""
        self.log_message.emit(f"💡 如果下载失败, 请优先检查网络, 需要科学上网")
        for attempt in range(1, MAX_RETRIES + 1):
            if self.is_cancelled:
                return False

            try:

                self.log_message.emit(f"🎬 开始下载视频 (尝试 {attempt}/{MAX_RETRIES})")

                # 首先提取视频信息获取标题和分辨率信息
                if attempt == 1:  # 只在第一次尝试时提取信息
                    try:
                        # 再次检查取消状态
                        if self.is_cancelled:
                            return False

                        info_opts = {
                            'cookies': 'cookies.txt',
                            'quiet': True,
                        }
                        with yt_dlp.YoutubeDL(info_opts) as ydl:
                            info = ydl.extract_info(url, download=False)
                            video_title = info.get('title', '未知标题')
                            self.video_info_extracted.emit(video_title)
                            self.log_message.emit(f"📝 视频标题: {video_title}")

                            # 检查分辨率限制
                            if self.resolution and not self._check_resolution_available(info):
                                self.log_message.emit(f"❌ 找不到符合分辨率限制的视频格式 (要求: ≥{self.resolution}p)")
                                self.log_message.emit("💡 建议: 尝试选择'不限制'分辨率或降低分辨率要求")
                                return False

                    except Exception as e:
                        self.log_message.emit(f"💡 请点击 校验网络 按钮, 进行网络测试")
                        self.log_message.emit(f"⚠️ 获取视频信息失败: {e}")

                # 再次检查取消状态
                if self.is_cancelled:
                    return False

                with yt_dlp.YoutubeDL(self.get_ydl_opts(output_folder)) as ydl:
                    ydl.download([url])

                self.log_message.emit("✅ 视频下载成功并处理完成")
                return True

            except Exception as e:
                error_msg = str(e)

                # 检查是否是用户取消导致的异常
                if "下载已被用户取消" in error_msg or self.is_cancelled:
                    self.log_message.emit("🛑 下载已被用户取消")
                    return False

                self.log_message.emit(f"❌ 下载失败 (尝试 {attempt}/{MAX_RETRIES}): {error_msg}")

                if attempt < MAX_RETRIES:
                    self.log_message.emit(f"⏳ 等待 2 秒后重试...")
                    # 在等待期间也检查取消状态
                    for i in range(20):  # 分成20个0.1秒的检查
                        if self.is_cancelled:
                            return False
                        time.sleep(0.1)
                else:
                    self.log_message.emit(f"💡 请点击 网络校验按钮 , 检查网 ")
                    self.log_message.emit(f"💥 最终下载失败: {error_msg}")
                    return False

        return False

    def download_playlist(self, url):
        """下载播放列表"""
        try:
            self.log_message.emit("📋 检测到播放列表，正在提取信息...")
            self.status_updated.emit("提取播放列表信息中...")

            # 检查取消状态
            if self.is_cancelled:
                return False

            # 提取播放列表信息
            ydl_opts_info = {
                'cookies': 'cookies.txt',
                'extract_flat': True,
                'quiet': True,
            }

            with yt_dlp.YoutubeDL(ydl_opts_info) as ydl:
                info = ydl.extract_info(url, download=False)

            entries = info.get('entries', [])
            playlist_title = info.get('title', '播放列表')
            playlist_folder = self.sanitize_filename(playlist_title)

            self.log_message.emit(f"📁 播放列表: {playlist_title}")
            self.log_message.emit(f"📊 共 {len(entries)} 个视频")

            # 创建播放列表文件夹
            full_playlist_path = os.path.join(self.download_path, playlist_folder)
            if not os.path.exists(full_playlist_path):
                os.makedirs(full_playlist_path)
                self.log_message.emit(f"📂 创建文件夹: {playlist_folder}")

            # 设置总视频数
            self.total_videos = len(entries)
            self.downloaded_videos = 0

            # 下载所有视频
            for idx, entry in enumerate(entries, 1):
                if self.is_cancelled:
                    break

                video_url = entry.get('url') or f"https://www.youtube.com/watch?v={entry['id']}"
                video_title = entry.get('title', f'视频 {idx}')

                self.log_message.emit(f"📹 下载第 {idx}/{len(entries)} 个视频: {video_title}")

                success = self.download_single_video(video_url, full_playlist_path)
                if not success and not self.is_cancelled:
                    self.log_message.emit(f"⚠️ 第 {idx} 个视频下载失败，继续下载下一个")

            return not self.is_cancelled

        except Exception as e:
            self.log_message.emit(f"💥 播放列表下载失败: {str(e)}")
            return False

    def run(self):
        """线程主函数"""
        try:
            self.log_message.emit(f"🚀 开始下载任务")
            self.log_message.emit(f"🔗 URL: {self.url}")
            self.log_message.emit(f"📁 下载路径: {self.download_path}")

            if self.resolution:
                self.log_message.emit(f"🎥 分辨率限制: {self.resolution}p")
            if self.only_mp4:
                self.log_message.emit(f"📹 格式限制: 仅MP4")

            # 检查是否为播放列表
            playlist_match = re.search(r"list=([A-Za-z0-9_-]+)", self.url)

            if playlist_match:
                success = self.download_playlist(self.url)
            else:
                self.total_videos = 1
                self.downloaded_videos = 0
                success = self.download_single_video(self.url, self.download_path)

            # 发送完成信号
            if self.is_cancelled:
                self.download_finished.emit(False, "下载已取消")
                self.log_message.emit("🛑 下载已取消")
            elif success:
                self.download_finished.emit(True, "下载完成")
                self.log_message.emit("🎉 所有下载任务完成")
                self.progress_updated.emit(100)
            else:
                self.download_finished.emit(False, "下载失败")
                self.log_message.emit("💥 下载任务失败")

        except Exception as e:
            error_msg = f"下载过程中发生错误: {str(e)}"
            self.log_message.emit(f"💥 {error_msg}")
            self.download_finished.emit(False, error_msg)

    def cancel_download(self):
        """取消下载"""
        self.is_cancelled = True
        self.log_message.emit("🛑 正在取消下载...")

        # 尝试强制终止线程（如果需要的话）
        if self.isRunning():
            self.requestInterruption()

